/********************************************************************
filename:   MP4Encoder.cpp
created:    2013-04-16
author:     firehood
purpose:    MP4�����������ڿ�Դ��mp4v2ʵ�֣�https://code.google.com/p/mp4v2/����
*********************************************************************/
#include "MP4Encoder.h"
#include <string.h>
#include "assert.h"

#ifndef _WIN32
#include "netinet/in.h"
#endif
#ifdef _SUPPORT_MP4
#define FRAME_BUF_SIZE (1000 << 10)


#define MIN_FRAME_SIZE 32
#define VIDEO_TIME_SCALE 90000
#define AUDIO_TIME_SCALE 8000
#define MOVIE_TIME_SCALE VIDEO_TIME_SCALE
#define PTS2TIME_SCALE(CurPTS, PrevPTS, timeScale) \
	((MP4Duration)((CurPTS - PrevPTS) * 1.0 / (double)(1e+3) * timeScale))
#define INVALID_PTS 0xFFFFFFFFFFFFFFFF
#define DEFAULT_VIDEO_TRACK_NUM 3
#define DEFAULT_VIDEO_PROFILE_LEVEL 1
#define DEFAULT_AUDIO_PROFILE_LEVEL 2

enum H264NALTYPE
{
	H264NT_NAL = 0,
	H264NT_SLICE,        //P 帧
	H264NT_SLICE_DPA,
	H264NT_SLICE_DPB,
	H264NT_SLICE_DPC,
	H264NT_SLICE_IDR, // I 帧
	H264NT_SEI,
	H264NT_SPS,
	H264NT_PPS,
};

CMP4Encoder::CMP4Encoder(void):
	m_nWidth(960),
	m_nHeight(576),
	m_nFrameRate(0)
{
	m_audioId     = NULL;
  	m_dataBuf     = new unsigned char[FRAME_BUF_SIZE];
	m_videoTrack  = MP4_INVALID_TRACK_ID;
	m_jpegVideoID = MP4_INVALID_TRACK_ID;
	m_bSetPps     = false;
	m_nTimeScale  = 0;
	m_bFirstAudio =  true;
	m_bFirstVideo =  true;
    m_u64AudioPTS =  0;
    m_u64VideoPTS =  0;
	m_u64FirstPTS =  INVALID_PTS;
	m_u64LastPTS  =  0;
}

CMP4Encoder::~CMP4Encoder(void)
{
	delete [] m_dataBuf;
}


MP4FileHandle CMP4Encoder::CreateMP4File(const char *pFileName,int width,int height,int timeScale/* = 90000*/,int frameRate/* = 30*/)
{
	if(NULL == pFileName )
	{
		printf("ERROR:No File name.\n");
		return NULL;
	}
	// create mp4 file
	MP4FileHandle hMp4file = MP4Create(pFileName);
	if (MP4_INVALID_FILE_HANDLE == hMp4file)
	{
		printf("ERROR:Open file failed. err %d .\n",errno);

		return NULL;
	}
	m_nWidth	 = width;
	m_nHeight	 = height;
	m_nTimeScale = timeScale;
	m_nFrameRate = frameRate;//30; //frameRate;
	m_hMp4file   = hMp4file;
	MP4SetTimeScale(m_hMp4file, m_nTimeScale);
	m_videoTrack    = MP4_INVALID_TRACK_ID;
	m_jpegVideoID= MP4_INVALID_TRACK_ID;
	m_audioId  = MP4_INVALID_TRACK_ID;
	m_bSetPps    = false;
	TRACE("Create MP4 file timeScale %d frame Rate %d .\n",
		             timeScale,
					 frameRate);
	m_uSecond     = 0;
	m_bFirstAudio =  true;
	m_bFirstVideo =  true;
	m_u64AudioPTS =  0;
	m_u64VideoPTS =  0;
	m_u64FirstPTS =  INVALID_PTS;
	m_u64LastPTS  =  0;
	return hMp4file;
}

bool CMP4Encoder::Write264Metadata(LPMP4ENC_Metadata lpMetadata,bool bAddAudio)
{
	m_videoTrack = MP4AddH264VideoTrack
		(m_hMp4file,
		m_nTimeScale,
		m_nTimeScale / m_nFrameRate,
		m_nWidth, // width
		m_nHeight,// height
		lpMetadata->Sps[1], // sps[1] AVCProfileIndication
		lpMetadata->Sps[2], // sps[2] profile_compat
		lpMetadata->Sps[3], // sps[3] AVCLevelIndication
		3);           // 4 bytes length before each NAL unit
	if (MP4_INVALID_TRACK_ID == m_videoTrack)
	{
		printf("add video track failed.\n");
		return false;
	}
	MP4SetVideoProfileLevel(m_hMp4file, 0x7F); //  Simple Profile @ Level 3

	// write sps
	MP4AddH264SequenceParameterSet(m_hMp4file,m_videoTrack,lpMetadata->Sps,lpMetadata->nSpsLen);
	// write pps
    MP4AddH264PictureParameterSet(m_hMp4file,m_videoTrack,lpMetadata->Pps,lpMetadata->nPpsLen);
	if(bAddAudio)
	{
		m_audioId = MP4AddAudioTrack(m_hMp4file, 48000, 1024, MP4_MPEG4_AUDIO_TYPE);
		if (MP4_INVALID_TRACK_ID == m_audioId )
		{
			printf("add audio track failed.\n");
			return false;
		}
		MP4SetAudioProfileLevel(m_hMp4file, 0x2);
	}
	return true;
}
int CMP4Encoder::WriteVideoData(const unsigned char * pData,int size, uint64_t u64PTS)
{
	if(NULL == pData)
	{
		return -1;
	}
	MP4Duration curDuration = 3600;
	uint64_t    diff        = 33;

	if ( 0 != u64PTS && m_bFirstVideo)
	{
		if (m_u64FirstPTS > u64PTS)
			m_u64FirstPTS = u64PTS;
		m_u64VideoPTS = u64PTS;
		m_bFirstVideo = false;
	}

	if(MP4_INVALID_TRACK_ID == m_jpegVideoID)
	{
		m_jpegVideoID = MP4AddVideoTrack(
						m_hMp4file,
						m_nTimeScale,
						m_nTimeScale / m_nFrameRate,
				        m_nWidth,
						m_nHeight,
			            MP4_JPEG_VIDEO_TYPE );
		if (MP4_INVALID_TRACK_ID == m_jpegVideoID)
		{
			TRACE("add video track failed.\n");
			return 0;
		}
	}

	diff = GetFrmDiffTm(u64PTS,m_u64VideoPTS);
	#define VIDEO_TIME_SCALE 90000
	curDuration = (MP4Duration)((diff*1000 * 1.0)/ (double)(1e+6) *VIDEO_TIME_SCALE);
	if(!MP4WriteSample(m_hMp4file, m_jpegVideoID, pData, size,curDuration, 0, 1))
	{
		TRACE("Write sample fail.\n");
		return -1;
	}
	m_u64LastPTS = m_u64VideoPTS = u64PTS;

	return size;
}

int CMP4Encoder::WriteH264VideoData(const unsigned char* pData,int size,uint64_t u64PTS)
{
	if(pData == NULL)
	{
		return -1;
	}
	MP4ENC_NaluUnit nalu;
	int pos = 0, len = 0;
	int datalen;
	MP4Duration curDuration = 3600;
	uint64_t    diff        = 33;

	if (size < MIN_FRAME_SIZE)
	{
		return MP4ENCODER_ENONE;
	}

	uchar_t * pDestBuf = NULL;
 	bool result = false;

	if (m_bFirstVideo)
	{
		if (m_u64FirstPTS > u64PTS)
			m_u64FirstPTS = u64PTS;
		m_u64VideoPTS = u64PTS;
		m_bFirstVideo = false;
	}

    while (len = ReadOneNaluFromBuf(pData,size,pos,nalu))
    {
		pos += len;
		if(nalu.type == 0x07 && MP4_INVALID_TRACK_ID == m_videoTrack ) // sps
		{
			// 添加h264 track
			m_videoTrack = MP4AddH264VideoTrack
				(m_hMp4file,
				m_nTimeScale,
				m_nTimeScale / m_nFrameRate,
				m_nWidth,     // width
				m_nHeight,    // height
				nalu.data[1], // sps[1] AVCProfileIndication
				nalu.data[2], // sps[2] profile_compat
				nalu.data[3], // sps[3] AVCLevelIndication
				3);           // 4 bytes length before each NAL unit
			if (MP4_INVALID_TRACK_ID == m_videoTrack)
			{
				TRACE_ERR("add video track failed.\n");
				return 0;
			}
			TRACE(" MP4AddH264VideoTrack m_nFrameRate %d m_nTimeScale %d m_nWidth %d m_nHeight %d ",
				     m_nFrameRate,m_nTimeScale,m_nWidth,m_nHeight);
			MP4SetVideoProfileLevel(m_hMp4file, 1); //  Simple Profile @ Level 3
			MP4AddH264SequenceParameterSet(m_hMp4file,m_videoTrack,nalu.data,nalu.size);

		}
		else if(nalu.type == 0x08 && !m_bSetPps) // pps
		{
			if(MP4_INVALID_TRACK_ID == m_videoTrack)
			{
				 return 0;
			}
			MP4AddH264PictureParameterSet(m_hMp4file,m_videoTrack,nalu.data,nalu.size);
			m_bSetPps = true;
		}
		else
        {
            if( 0 == m_videoTrack)
            {
                return 0;
            }

			if( nalu.size <  0 ||
				nalu.size >= FRAME_BUF_SIZE)
			{
				continue;
			}
			if( nalu.type  == 0x06 ||
				nalu.type  == 0x07 ||
				nalu.type == 0x08)
			{
				continue;
			}
            datalen = nalu.size+4;


			pDestBuf = (uchar_t*)(nalu.data - 4);
			pDestBuf[0] = nalu.size>>24;
			pDestBuf[1] = nalu.size>>16;
			pDestBuf[2] = nalu.size>>8;
			pDestBuf[3] = nalu.size&0xff;

			diff = GetFrmDiffTm(u64PTS,m_u64VideoPTS);
			curDuration = (MP4Duration)((diff *1000 * 1.0)/ (double)(1e+6) *VIDEO_TIME_SCALE);
			if (nalu.type == 5)
			{
				result = MP4WriteSample(m_hMp4file, m_videoTrack, pDestBuf, datalen,
					curDuration);
			}
			else
			{
				result = MP4WriteSample(m_hMp4file, m_videoTrack, pDestBuf, datalen,
					curDuration, 0, false);
			}

			m_u64LastPTS = m_u64VideoPTS = u64PTS;
        }
    }

	return pos;
}



int CMP4Encoder::WriteH264AudioData(const unsigned char* pData,int size,uint64_t u64PTS)
{
	if(MP4_INVALID_TRACK_ID == m_audioId)
	{
		m_audioId = MP4AddAudioTrack(
			m_hMp4file,
			m_nTimeScale,
			m_nTimeScale / m_nFrameRate,
			MP4_MPEG4_AUDIO_TYPE );

		if (MP4_INVALID_TRACK_ID == m_audioId )
		{
			TRACE("add video track failed.\n");
			return 0;
		}
	}


	if (m_bFirstAudio)
	{
		if (m_u64FirstPTS > u64PTS)
			m_u64FirstPTS = u64PTS;
		m_u64AudioPTS = u64PTS;
		m_bFirstAudio = false;
	}

	MP4Duration curDuration = 3600;
	uint64_t    diff        = 33;


	diff = GetFrmDiffTm(u64PTS,m_u64AudioPTS);
	curDuration = (MP4Duration)((diff *1000 * 1.0)/ (double)(1e+6) *VIDEO_TIME_SCALE);
	if(!MP4WriteSample(m_hMp4file, m_audioId, pData, size,curDuration, 0, 1))
	{
		return -1;
	}

	m_u64LastPTS = m_u64AudioPTS = u64PTS;
	return 0;
}

int CMP4Encoder::ReadOneNaluFromBuf(const unsigned char *buffer,unsigned int nBufferSize,unsigned int offSet,MP4ENC_NaluUnit &nalu)
{
	unsigned int i   = offSet;
	unsigned int pos = 0;

	while(i < nBufferSize)
	{
		if( buffer[i++] == 0x00 &&
			buffer[i++] == 0x00 &&
			buffer[i++] == 0x00 &&
			buffer[i++] == 0x01
			)
		{
			pos = i;
			while (pos < nBufferSize)
			{
				if( buffer[pos++] == 0x00 &&
					buffer[pos++] == 0x00 &&
					buffer[pos++] == 0x00 &&
					buffer[pos++] == 0x01
					)
				{
					break;
				}
			}
			if(pos == nBufferSize)
			{
				nalu.size = pos-i;
			}
			else
			{
				nalu.size = (pos-4)-i;
			}

			nalu.type = buffer[i]&0x1f;
			nalu.data =(unsigned char*)&buffer[i];
			return (nalu.size+i-offSet);
		}
	}

	return 0;
}

void CMP4Encoder::CloseMP4File()
{
	if(m_hMp4file)
	{
		MP4Close(m_hMp4file);
#ifndef _WIN32
		system("sync");
#endif
		m_hMp4file = NULL;
	}
}

bool CMP4Encoder:: PraseMetadata(const unsigned char* pData,int size,MP4ENC_Metadata &metadata)
{
	if(pData == NULL || size<4)
	{
		return false;
	}
	MP4ENC_NaluUnit nalu;
	int pos = 0;
	bool bRet1 = false;
	bool bRet2 = false;

	while (int len = ReadOneNaluFromBuf(pData,size,pos,nalu))
	{
		if(nalu.type == 0x07)
		{
			memcpy(metadata.Sps,nalu.data,nalu.size);
			metadata.nSpsLen = nalu.size;
			bRet1 = true;
		}
		else if((nalu.type == 0x08))
		{
			memcpy(metadata.Pps,nalu.data,nalu.size);
			metadata.nPpsLen = nalu.size;
			bRet2 = true;
		}
		pos += len;
	}
	if(bRet1 && bRet2)
	{
		return true;
	}
	return false;
}

MP4EncoderResult CMP4Encoder::MP4AddH264Track(const uchar_t *sData, int nSize,
											 int nWidth, int nHeight, int nFrameRate/* = 25 */)
{
	int sps, pps;
	for (sps = 0; sps < nSize;)
	{
		if (sData[sps++] == 0x00 && sData[sps++] == 0x00 && sData[sps++] == 0x00
			&& sData[sps++] == 0x01)
			break;
	}

	for (pps = sps; pps < nSize;)
	{
		if (sData[pps++] == 0x00 && sData[pps++] == 0x00 && sData[pps++] == 0x00
			&& sData[pps++] == 0x01)
			break;
	}
	if (sps >= nSize || pps >= nSize)
		return MP4ENCODER_ERROR(MP4ENCODER_E_ADD_VIDEO_TRACK);

	m_videoTrack = MP4AddH264VideoTrack(m_hMp4file, VIDEO_TIME_SCALE,
		VIDEO_TIME_SCALE / nFrameRate, nWidth, nHeight,
		sData[sps + 1], sData[sps + 2], sData[sps + 3], DEFAULT_VIDEO_TRACK_NUM);
	if (MP4_INVALID_TRACK_ID == m_videoTrack)
		return MP4ENCODER_ERROR(MP4ENCODER_E_ADD_VIDEO_TRACK);

	MP4SetVideoProfileLevel(m_hMp4file, DEFAULT_VIDEO_PROFILE_LEVEL);
	MP4AddH264SequenceParameterSet(m_hMp4file, m_videoTrack, sData + sps,
		pps - sps - 4);
 	MP4AddH264PictureParameterSet(m_hMp4file, m_videoTrack, sData + pps,
 		nSize - pps);

	return MP4ENCODER_ENONE;
}

uint64_t CMP4Encoder::GetFrmDiffTm(uint64_t curTime, uint64_t lastTime)
{
	uint64_t diff = 33;
	if(curTime >= lastTime)
	{
		diff =  curTime - lastTime;
		if(diff > 1000)
		{
			diff = 1000;//保证1秒种一定会有一帧
		}
	}
	else
	{
		diff = 33;
	}
	return diff;
}

SIZE CMP4Encoder::GetJPGImageSize(uchar_t* buffer,int len)
{
	SIZE sizePiex={0,0};
	do
	{
		if (buffer==NULL || len<=2)
		{
			CROSS_ASSERT(false);
		}

		uchar_t* cursor=buffer;

		//图像头SOI标识
		if(*cursor!=0xff || *(++cursor)!=0xD8)
		{
			CROSS_ASSERT(false);
			break;
		}

		bool bDone = false;

		while(! bDone)
		{
			_tagImgSection_t ImgSection;
			memset(&ImgSection, 0, sizeof(ImgSection));
			//读入段头（段标记及长度标记共4字节）
			ImgSection.SectionStart=*(++cursor);
			ImgSection.SectionID=*(++cursor);
			ImgSection.SectionLenH=*(++cursor);
			ImgSection.SectionLenL=*(++cursor);

			if(ImgSection.SectionStart != 0xFF) //非段标识标记
				break;

			switch(ImgSection.SectionID)
			{
			case(0xFF): //当前标识为数据
				{
					continue;
				}
			case(0xD9): // EOI，End of Image，图像结束
				{
					bDone = TRUE;
					break;
				}
			case(0xC0): //SOF0，Start of Frame，帧图像开始
				{

					SOF0STRUCT sofoData = {0};
					sofoData.Head.Sample_precision=*(++cursor);
					sofoData.Head.HighPiexsH=*(++cursor);
					sofoData.Head.HighPiexsL=*(++cursor);
					sofoData.Head.WidthPiexsH=*(++cursor);
					sofoData.Head.WidthPiexsL=*(++cursor);
					sizePiex.cx = MAKEWORD(sofoData.Head.WidthPiexsL, sofoData.Head.WidthPiexsH);
					sizePiex.cy = MAKEWORD(sofoData.Head.HighPiexsL, sofoData.Head.HighPiexsH);
					bDone = TRUE;

					break;
				}
			default:
				{
					int iSeek = (int)MAKEWORD(ImgSection.SectionLenL, ImgSection.SectionLenH) - 2;
					cursor+=iSeek;
					break;
				}
			}
		}

	} while (0);

	return sizePiex;
}

#endif

